#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF DOI提取器 - 使用QWen2.5-vl API
第一步：检测标题
第二步：通过QWen2.5-vl API解析PDF获取DOI
"""

import fitz  # PyMuPDF
import re
import os
import base64
import requests
import json
from typing import Dict, List, Optional, Tuple
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class PDFDOIExtractor:
    """PDF DOI提取器 - 使用QWen2.5-vl API"""
    
    def __init__(self, pdf_dir: str, api_key: str = None, api_base: str = None):
        """
        初始化PDF DOI提取器
        
        Args:
            pdf_dir: PDF文件所在目录路径
            api_key: QWen2.5-vl API密钥
            api_base: QWen2.5-vl API基础URL
        """
        self.pdf_dir = pdf_dir
        
        # QWen2.5-vl API配置
        self.api_key = api_key or os.getenv('QWEN_API_KEY', 'your-api-key-here')
        self.api_base = api_base or os.getenv('QWEN_API_BASE', 'https://api.qwen.ai/v1')
        
        self.pdf_files = []
        self._scan_pdf_files()
    
    def _scan_pdf_files(self):
        """扫描目录下的所有PDF文件"""
        if not os.path.exists(self.pdf_dir):
            raise FileNotFoundError(f"目录不存在: {self.pdf_dir}")
        
        for file in os.listdir(self.pdf_dir):
            if file.lower().endswith('.pdf'):
                self.pdf_files.append(os.path.join(self.pdf_dir, file))
        
        logger.info(f"找到 {len(self.pdf_files)} 个PDF文件")
    

    

    

    

    
    def extract_title_doi_author_single_request(self, pdf_path: str) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        使用一次QWen2.5-vl API调用同时获取标题、DOI和作者
        
        Args:
            pdf_path: PDF文件路径
            
        Returns:
            (标题, DOI, 作者)的元组
        """
        try:
            logger.info(f"🔍 使用一次API调用同时获取标题和DOI: {os.path.basename(pdf_path)}")
            
            # 只处理第一页
            doc = fitz.open(pdf_path)
            if doc.page_count == 0:
                logger.warning("PDF文件没有页面")
                return None, None
            
            # 转换第一页为图像
            page = doc.load_page(0)
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
            img_data = pix.tobytes("png")
            doc.close()
            
            # 使用QWen2.5-vl API同时提取标题、DOI和作者
            title, doi, author = self._extract_title_doi_author_with_qwen_api(img_data, 0)
            return title, doi, author
            
        except Exception as e:
            logger.error(f"标题、DOI和作者提取失败: {e}")
            return None, None, None
    
    def count_complete_records(self, results: List[Dict[str, Optional[str]]]) -> int:
        """
        统计包含完整信息（标题、DOI、作者）的记录数量
        
        Args:
            results: 处理结果列表
            
        Returns:
            完整信息记录数量
        """
        complete_count = 0
        for result in results:
            if result['title'] and result['doi'] and result['author']:
                complete_count += 1
        return complete_count
    
    def _extract_title_doi_author_with_qwen_api(self, image_data: bytes, page_num: int) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        使用QWen2.5-vl API从图像中同时提取标题、DOI和作者
        
        Args:
            image_data: 图像数据（PNG格式）
            page_num: 页码
            
        Returns:
            (标题, DOI, 作者)的元组
        """
        try:
            # 将图像数据转换为base64
            image_base64 = base64.b64encode(image_data).decode('utf-8')
            
            # 构建API请求
            headers = {
                'Authorization': f'Bearer {self.api_key}',
                'Content-Type': 'application/json'
            }
            
            # 构建提示词，同时要求提取标题、DOI和作者
            prompt = f"""请仔细查看这个PDF页面的图像，找到其中的论文标题、DOI（数字对象唯一标识符）和作者信息。

请按照以下格式返回结果：
标题：[论文标题]
DOI：[DOI内容]
作者：[作者信息]

说明：
1. 如果找到标题，请直接填写标题内容
2. 如果找到DOI，请直接填写DOI内容，格式如：10.xxxx/xxxxx
3. 如果找到作者，请直接填写作者姓名或作者列表
4. 如果没有找到标题，请填写：未找到标题
5. 如果没有找到DOI，请填写：未找到DOI
6. 如果没有找到作者，请填写：未找到作者
7. 如果图像不清晰无法识别，请填写：图像不清晰

请严格按照上述格式返回，不要添加其他解释。"""
            
            payload = {
                "model": "qwen2.5-vl-32b-instruct",
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{image_base64}"
                                }
                            }
                        ]
                    }
                ],
                "max_tokens": 300,
                "temperature": 0.1
            }
            
            logger.info(f"   正在调用QWen2.5-vl API同时提取标题和DOI...")
            
            # 发送API请求
            session = requests.Session()
            session.mount('https://', requests.adapters.HTTPAdapter(
                max_retries=3,
                pool_connections=10,
                pool_maxsize=10
            ))
            
            response = session.post(
                f"{self.api_base}/chat/completions",
                headers=headers,
                json=payload,
                timeout=60,
                verify=True
            )
            
            if response.status_code == 200:
                result = response.json()
                content = result['choices'][0]['message']['content'].strip()
                
                logger.info(f"   QWen2.5-vl API响应: {content}")
                
                # 解析标题、DOI和作者
                title = None
                doi = None
                author = None
                
                # 提取标题
                title_match = re.search(r'标题：(.+?)(?:\n|$)', content)
                if title_match:
                    title = title_match.group(1).strip()
                    if title == "未找到标题":
                        title = None
                
                # 提取DOI
                doi_match = re.search(r'DOI：(.+?)(?:\n|$)', content)
                if doi_match:
                    doi = doi_match.group(1).strip()
                    if doi == "未找到DOI":
                        doi = None
                
                # 提取作者
                author_match = re.search(r'作者：(.+?)(?:\n|$)', content)
                if author_match:
                    author = author_match.group(1).strip()
                    if author == "未找到作者":
                        author = None
                
                logger.info(f"✅ 在第{page_num}页解析结果 - 标题: {title}, DOI: {doi}, 作者: {author}")
                return title, doi, author
            else:
                logger.error(f"   QWen2.5-vl API调用失败: {response.status_code} - {response.text}")
                return None, None
                
        except Exception as e:
            logger.error(f"   QWen2.5-vl API调用异常: {e}")
            return None, None
    
    def process_single_pdf(self, pdf_path: str) -> Dict[str, Optional[str]]:
        """
        处理单个PDF文件：检测标题和提取DOI
        
        Args:
            pdf_path: PDF文件路径
            
        Returns:
            包含标题和DOI的字典
        """
        result = {
            'file_path': pdf_path,
            'file_name': os.path.basename(pdf_path),
            'title': None,
            'doi': None,
            'author': None
        }
        
        try:
            # 使用一次API调用同时获取标题、DOI和作者
            title, doi, author = self.extract_title_doi_author_single_request(pdf_path)
            result['title'] = title
            result['doi'] = doi
            result['author'] = author
            
            # 输出结果
            logger.info(f"\n📊 处理结果: {os.path.basename(pdf_path)}")
            logger.info(f"   标题: {title if title else '未检测到'}")
            logger.info(f"   DOI: {doi if doi else '未找到'}")
            logger.info(f"   作者: {author if author else '未找到'}")
            
        except Exception as e:
            logger.error(f"处理PDF文件失败: {e}")
        
        return result
    
    def batch_process(self, max_files: int = None, target_complete_count: int = 100) -> List[Dict[str, Optional[str]]]:
        """
        批量处理PDF文件，当找到指定数量的完整信息时自动停止
        
        Args:
            max_files: 最大处理文件数量，None表示处理所有文件
            target_complete_count: 目标完整信息数量，达到后自动停止
            
        Returns:
            处理结果列表
        """
        if max_files is None:
            max_files = len(self.pdf_files)
        
        results = []
        total_files = min(max_files, len(self.pdf_files))
        complete_count = 0
        
        logger.info(f"🚀 开始批量处理，共 {total_files} 个文件")
        logger.info(f"🎯 目标：找到 {target_complete_count} 个包含完整信息（标题、DOI、作者）的数据")
        
        for i, pdf_path in enumerate(self.pdf_files[:max_files]):
            logger.info(f"\n{'='*60}")
            logger.info(f"📄 处理文件 {i+1}/{total_files}: {os.path.basename(pdf_path)}")
            logger.info(f"{'='*60}")
            
            result = self.process_single_pdf(pdf_path)
            results.append(result)
            
            # 检查是否包含完整信息
            if result['title'] and result['doi'] and result['author']:
                complete_count += 1
                logger.info(f"🎉 找到第 {complete_count} 个完整信息！")
                
                # 达到目标数量时停止
                if complete_count >= target_complete_count:
                    logger.info(f"🎯 已达到目标数量 {target_complete_count}，自动停止处理！")
                    break
        
        # 输出汇总
        logger.info(f"\n📊 批量处理完成！")
        logger.info(f"   总文件数: {len(results)}")
        logger.info(f"   成功提取标题: {sum(1 for r in results if r['title'])}")
        logger.info(f"   成功提取DOI: {sum(1 for r in results if r['doi'])}")
        logger.info(f"   成功提取作者: {sum(1 for r in results if r['author'])}")
        logger.info(f"   完整信息数量: {complete_count}")
        
        # 第三步：保存结果到文件
        self.save_results_to_file(results)
        
        return results
    
    def save_results_to_file(self, results: List[Dict[str, Optional[str]]], output_file: str = None):
        """
        将处理结果保存到文件
        
        Args:
            results: 处理结果列表
            output_file: 输出文件名，默认为results_YYYYMMDD_HHMMSS.csv
        """
        try:
            import csv
            from datetime import datetime
            
            # 生成默认文件名
            if output_file is None:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                output_file = f"results_{timestamp}.csv"
            
            # 确保输出目录存在
            output_dir = os.path.dirname(output_file) if os.path.dirname(output_file) else "."
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            # 保存为CSV文件
            with open(output_file, 'w', newline='', encoding='utf-8-sig') as csvfile:
                fieldnames = ['文件路径', '文件名', '标题', 'DOI', '作者', '处理时间']
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                
                writer.writeheader()
                
                for result in results:
                    writer.writerow({
                        '文件路径': result['file_path'],
                        '文件名': result['file_name'],
                        '标题': result['title'] if result['title'] else '未找到',
                        'DOI': result['doi'] if result['doi'] else '未找到',
                        '作者': result['author'] if result['author'] else '未找到',
                        '处理时间': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
            
            logger.info(f"✅ 结果已保存到: {output_file}")
            
            # 同时保存为JSON文件（便于程序读取）
            json_file = output_file.replace('.csv', '.json')
            with open(json_file, 'w', encoding='utf-8') as jsonfile:
                json.dump(results, jsonfile, ensure_ascii=False, indent=2)
            
            logger.info(f"✅ 结果已同时保存到: {json_file}")
            
            # 输出统计信息
            total_files = len(results)
            title_success = sum(1 for r in results if r['title'])
            doi_success = sum(1 for r in results if r['doi'])
            author_success = sum(1 for r in results if r['author'])
            complete_count = self.count_complete_records(results)
            
            logger.info(f"\n📈 保存结果统计:")
            logger.info(f"   总文件数: {total_files}")
            logger.info(f"   标题提取成功率: {title_success}/{total_files} ({title_success/total_files*100:.1f}%)")
            logger.info(f"   DOI提取成功率: {doi_success}/{total_files} ({doi_success/total_files*100:.1f}%)")
            logger.info(f"   作者提取成功率: {author_success}/{total_files} ({author_success/total_files*100:.1f}%)")
            logger.info(f"   完整信息数量: {complete_count}/{total_files} ({complete_count/total_files*100:.1f}%)")
            
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
            import traceback
            logger.error(f"错误详情: {traceback.format_exc()}")


def test_single_pdf():
    """测试单个PDF文件"""
    try:
        extractor = PDFDOIExtractor("/Users/hutao/Desktop/8.22/")
        
        # 测试第一个PDF文件
        if extractor.pdf_files:
            test_file = extractor.pdf_files[0]
            logger.info(f"测试文件: {os.path.basename(test_file)}")
            
            result = extractor.process_single_pdf(test_file)
            
            logger.info(f"\n🎯 测试结果:")
            logger.info(f"   文件: {result['file_name']}")
            logger.info(f"   标题: {result['title']}")
            logger.info(f"   DOI: {result['doi']}")
            
        else:
            logger.info("没有找到PDF文件")
            
    except Exception as e:
        logger.error(f"测试失败: {e}")


if __name__ == "__main__":
    test_single_pdf()
